#pragma once

#include <chrono>
#include <memory>
#include <sstream>
#include <string>
#include <vector>

#include <opencv2/opencv.hpp>

#include "tensorrt_onnx.h"

struct Object {
  cv::Rect2f box;
  float confidence;
  int label;
};

class YoloxDetector {
public:
  explicit YoloxDetector(const std::string &model_path)
      : model_path_(model_path), tensorrt_onnx_(nullptr), model_height_(640),
        model_width_(640) {}

  virtual ~YoloxDetector() = default;

  bool Init();

  bool Detect(cv::Mat &input_image, std::vector<Object> *objects);

private:
  void DecodeYoloxOutputs(const float *const output, const float width_scale,
                          const float height_scale,
                          const float confidence_thresh,
                          std::vector<Object> *objs);

private:
  std::string model_path_;
  std::unique_ptr<TensorrtOnnxInference> tensorrt_onnx_;
  int model_height_;
  int model_width_;
};
